import os
import numpy as np
import pandas as pd
from loguru import logger

from mylib import gen_excel
from models.stock_model import StockNumber, DayInfo


def get_all_stock(today_str, N, stock):
    link_name_arr = []
    industry_arr = []
    price_arr = []
    over98_arr = []
    over7_arr = []
    over3_arr = []
    down3_arr = []
    up_avg_arr = []
    down_avg_arr = []
    cal_days_arr = []

    df = pd.read_csv('../../all.csv')
    count = 0
    for row in df.index:
        logger.info(f'row = {row}')
        # if row > 100:
        #     break
        if stock and df.loc[row]['ts_code'] != stock:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        if 'ST' in sn.name:
            continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('00'):
            continue
        count += 1
        logger.info(f'{sn.ts_code},{sn.name}')
        try:
            link_name, price, over98, over7, over3, down3, up_avg, down_avg, cal_days = analysis_stock(N, sn)
            if link_name is None:
                logger.warning(f'continue {sn.name}')
                continue
            logger.warning(
                f'{link_name},{sn.industry}, {price}, {over98}, {over7}, {over3}, {down3}, {up_avg}, {down_avg}, {cal_days}')
            link_name_arr.append(link_name)
            industry_arr.append(sn.industry)
            price_arr.append(price)
            over98_arr.append(over98)
            over7_arr.append(over7)
            over3_arr.append(over3)
            down3_arr.append(down3)
            up_avg_arr.append(up_avg)
            down_avg_arr.append(down_avg)
            cal_days_arr.append(cal_days)
        except Exception as e:
            logger.error(e)
            logger.error(f'continue {sn.name}')
    rank_dict = {
        'name': link_name_arr,
        'industry': industry_arr,
        'price': price_arr,
        'over98': over98_arr,
        'over7': over7_arr,
        'over3': over3_arr,
        'down3': down3_arr,
        'up_avg': up_avg_arr,
        'down_avg': down_avg_arr,
        'cal_days': cal_days_arr
    }
    file_name = f'cal_rank_{N}_{today_str}.xlsx'
    file_path = f'{log_dir}/{file_name}'
    if os.path.exists(file_path):
        os.remove(file_path)
    gen_excel.export_data(p_data=rank_dict,
                          p_file_path=file_path)


def get_top(socket_arr, sort_arr, top_n=10):
    ziparr = list(zip(socket_arr, sort_arr))
    sorted_ziparr = sorted(ziparr, key=lambda x: x[1])
    return set([item[0] for item in sorted_ziparr[:top_n]])


def analysis_stock(N, sn):
    csv_path = f'stocks/{sn.ts_code}.csv'
    df2 = pd.read_csv(csv_path)
    q10 = []
    q11 = []
    all_di = []
    parr = 0
    today_di = None
    cnt = 0
    pct_arr = []
    arr_high = []
    arr_low = []
    date_arr = []
    price_arr = []
    delta_days = []
    delta_d = 0
    price_arr_high = []
    price_arr_low = []

    up_avg_arr, down_avg_arr = [], []
    over98, over7, over3, down3, cal_days = 0, 0, 0, 0, 0
    for row2 in df2.index:
        cal_days += 1
        delta_d += 1
        di = DayInfo(sn, df2.loc[row2])
        di.name = sn.name
        di.ts_code = sn.ts_code
        di.industry = sn.industry
        # di.trade_date = df2.loc[row2]['trade_date']
        # di.open = df2.loc[row2]['open']
        # di.high = df2.loc[row2]['high']
        price_arr_high.append(di.high)
        # di.low = df2.loc[row2]['low']
        price_arr_low.append(di.low)
        # di.close = float(df2.loc[row2]['close'])
        if today_di and today_di.close < 10:
            return None, today_di.close, over98, over7, over3, down3, \
                   round(np.average(up_avg_arr), 2), round(np.average(down_avg_arr), 2), cal_days
        if cnt > 40:
            link_code_arr = sn.ts_code.split('.')
            link_code = f'{link_code_arr[1]}{link_code_arr[0]}'
            hyperlink = f'"https://xueqiu.com/S/{link_code}"'
            link_name = f'= HYPERLINK({hyperlink},"{sn.ts_code}_{sn.name}")'
            return link_name, today_di.close, over98, over7, over3, down3, \
                   round(np.average(up_avg_arr), 2), round(np.average(down_avg_arr), 2), cal_days
        # di.pre_close = df2.loc[row2]['pre_close']
        # di.change = df2.loc[row2]['change']
        # di.pct_chg = float(df2.loc[row2]['pct_chg'])
        if di.pct_chg >= 9.8:
            over98 += 1
        elif di.pct_chg >= 7:
            over7 += 1
        elif di.pct_chg >= 3:
            over3 += 1
        elif di.pct_chg <= -3:
            down3 += 1

        if di.pct_chg > 0:
            up_avg_arr.append(di.pct_chg)
        if di.pct_chg < 0:
            down_avg_arr.append(di.pct_chg)

        # di.vol = df2.loc[row2]['vol']
        # di.amount = df2.loc[row2]['amount']
        if len(q10) == N * 2 + 1:
            ad = all_di[N]
            if max(q10) == q10[N]:
                delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                delta_d = 0
                pct = round((parr - q10[N]) / q10[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    arr_high.append(pct)
                else:
                    arr_low.append(pct)
                price_arr.append(q10[N])
                date_arr.append(ad.trade_date)
                parr = q10[N]
                cnt += 1
            elif min(q11) == q11[N]:
                pct = round((parr - q11[N]) / q11[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    delta_days.append(delta_d) if delta_days else delta_days.append(delta_d - 6)
                    arr_high.append(pct)
                else:
                    delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                    arr_low.append(pct)
                delta_d = 0
                price_arr.append(q11[N])
                date_arr.append(ad.trade_date)
                parr = q11[N]
                cnt += 1
            all_di.pop(0)
            q10.pop(0)
            q11.pop(0)
        all_di.append(di)
        q10.append(di.high)
        q11.append(di.low)
        if today_di is None:
            parr = di.low
            today_di = di
            price_arr.append(di.low)
            date_arr.append(di.trade_date)


def printNumber(today_str, stocks, N):
    for stock in stocks:
        get_all_stock(today_str, N, stock)


if __name__ == '__main__':
    bkd = 0
    N = 6


    def msg_filter(msg):
        def is_msg(record):
            return msg in record['message']

        return is_msg


    stocks = [
        ''
    ]
    log_dir = os.path.abspath(__file__).replace('.py', '')
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)

    fpath = f'{log_dir}/rank.txt'
    if os.path.exists(fpath):
        os.remove(fpath)
    logger.add(fpath)
    logger.info('生成csv start')

    df = pd.read_csv('../../stocks/000001.SZ.csv')
    all_trade_date_arr = []
    for row in df.index:
        all_trade_date_arr.append(df.loc[row]['trade_date'])

    t_arr = []
    for idx, trade_date in enumerate(all_trade_date_arr):
        if idx > bkd:
            break
        today_str = str(trade_date)
        try:
            printNumber(today_str, stocks, N)
        except Exception as e:
            logger.error(e)
    logger.info('生成excel运行完成')
